import re
from io import StringIO
from string import ascii_lowercase, digits
lower_digits = ascii_lowercase + digits

import tybc_compile
from tybc_compile import CompileError

class LexError(Exception):
  pass

Keywords = {
  'function': 'FUNCTION',
  'write': 'WRITE',
  'read': 'READ',
  'while': 'WHILE',
  'if': 'IF',
  'else': 'ELSE',
  'elf': 'ELSIF',
  'end': 'END',
  'array': 'ARRAY',
  'et': 'ET',
  'ou': 'OU',
  'non': 'NON',
  'goto': 'GOTO',
  'label': 'LABEL',
  'vrai': 'VRAI',
  'faux': 'FAUX',
  'exit': 'EXIT',
  }

SYM2 = {
  ':=': 'DEF',
  '><': 'NE',
  '>=': 'GE',
  '<=': 'LE',
  }

SYM1 = {
  '=': 'EQ',
  '>': 'GT',
  '<': 'LT',
  '+': 'PLUS',
  '-': 'MINUS',
  '*': 'MUL',
  '/': 'DIV',
  '%': 'MOD',
  ';': 'SEMI',
  ',': 'COMMA',
  '$': 'DOLLAR',
  '(': 'LPAR',
  ')': 'RPAR',
  '[': 'LBRA',
  ']': 'RBRA',
  '|': 'PIPE',
  '{': 'LBLO',
  '}': 'RBLO',
  }

def lex(text):
  tok=[]
  line=0
  sio=StringIO(text)

  for s in sio:
    line+=1
    a=[]
    while True:
      s=s.lstrip()
      if len(s)==0 or s[0]=='#':
        break
      elif s[0] in ascii_lowercase:
        i=1
        while i<len(s) and s[i] in lower_digits:
          i+=1
        g=s[:i]
        if g in Keywords:
          a.append((Keywords[g], None, line))
        else:
          a.append(('ID', g, line))
        s=s[i:]
      elif s[0] in digits:
        i=1
        while i<len(s) and s[i] in digits:
          i+=1
        a.append(('NUM', int(s[:i]), line))
        s=s[i:]
      elif s[:2] in SYM2:
        a.append((SYM2[s[:2]], None, line))
        s=s[2:]
      elif s[0] in SYM1:
        a.append((SYM1[s[0]], None, line))
        s=s[1:]
      else:
        raise LexError('unknown token at "{0}" (line {1})'.format(s.rstrip(), line))
    if a:
      tok.extend(a)
      tok.append(('NL', None, line))

  sio.close()
  return tok

class AsmError(Exception):
  pass

def check(code, mac):
  labels=set()
  for i, a in enumerate(code):
    if isinstance(a, (int, str)):
      pass
    elif isinstance(a, tuple) and len(a)==2 and \
         isinstance(a[0], str) and isinstance(a[1], (int, str)):
      a0 = a[0]
      a1 = a[1]
      if a0=='lab':
        if not (isinstance(a1, str) and re.match('L\S*\Z', a1)):
          raise AsmError("label name error: "+a1)
        if a1 in labels:
          raise AsmError("label {0} is already appeared".format(a1))
        labels.add(a1)
    else:
      raise AsmError('code is invalid: ({0}, {1})'.format(a, i))
  for k, v in mac.items():
    if not (isinstance(k, str) and re.match('[A-KM-Z]\S*\Z', k)):
      raise AsmError('macro name error ({0}: {1})'.format(k, v))
    if not isinstance(v, int):
      raise AsmError('macro value must be int ({0}: {1})'.format(k, v))

def optimize(code):
  while True:
    ln=len(code)

    rep={}
    a=[]
    l=None
    for i in code:
      if isinstance(i, tuple):
        assert i[0]=='lab'
        if l is None:
          a.append(i)
          l=i[1]
        else:
          rep[i[1]]=l
      else:
        a.append(i)
        l=None
    code=[rep.get(i, i) for i in a]

    for i in range(1, len(code)):
      if (code[i-1], code[i]) in ((0, 'add'), (0, 'sub'), (1, 'mul'), (1,'div')):
        code[i-1]=code[i]=None
    code=[i for i in code if i is not None]

    for i in range(3, len(code)):
      if code[i-3]==1 and ('lab', code[i-2])==code[i] and code[i-1]=='jpo':
        code[i-3]=code[i-2]=code[i-1]=None
    code=[i for i in code if i is not None]

    if ln==len(code):
      return code

def asm(code, mac):
  check(code, mac)

  out=[]
  for u in code:
    if not isinstance(u, tuple):
      out.append(u)
      continue
    u0 = u[0]
    u1 = u[1]
    if u0=='sal':
      p = mac.get(u1, u1)
      if not (isinstance(p, int) and p>=0):
        raise AsmError('cannot expand {0}'.format(u))
      out.extend([0]*p)
    elif u0=='pop':
      p = mac.get(u1, u1)
      if not (isinstance(p, int) and p>=0):
        raise AsmError('cannot expand {0}'.format(u))
      out.extend(['pop']*p)
    elif u0=='god' or u0=='gto' or u0=='jpo' or u0=='hal':
      out.append(u1)
      out.append(u0)
    elif u0=='jmp':
      out.append(1)
      out.append(u1)
      out.append('jpo')
    elif u0=='lab':
      out.append(u)
    else:
      raise AsmError('cannot expand {0}'.format(u))

  out=[mac.get(i, i) for i in out]

  out=optimize(out)

  lab={}
  a=[]
  for c in out:
    if isinstance(c, tuple):
      assert c[0]=='lab'
      lab[c[1]]=len(a)
    else:
      a.append(c)
  out=[lab.get(i, i) for i in a]

  return out

def compile_allinone(txt, achk=True, tape=False):
  tok = lex(txt)
  code, mac = tybc_compile.compile(tok, achk, tape)
  prog = asm(code, mac)
  return prog

if __name__ == "__main__":
  import sys
  from optparse import OptionParser

  usage='usage: python %prog source [options]'
  op=OptionParser(usage)
  op.set_defaults(fout='epro.code')
  op.add_option('-o', dest='fout', metavar='file', help='output file')
  op.add_option('-a', default=True, action='store_false', dest='achk', help='not check array index')
  op.add_option('-t', default=False, action='store_true', dest='tape', help='use tape')
  opts, args = op.parse_args()
  if len(args)!=1:
    op.error('too many or few arguments')

  try:
    with open(args[0], encoding='latin-1') as f:
      t=f.read()
    prog=compile_allinone(t, opts.achk, opts.tape)
    with open(opts.fout, 'w') as f:
      for l in prog:
        f.write(str(l)+'\n')
  except (LexError, CompileError, IOError) as e: # AsmError is unexpected without compiler's bug
    sys.exit(e)
